import cv2
import numpy as np
import os
import yaml
from datetime import datetime

# 设置路径和参数
chessboard_path = "chessBoard1280"  # 保存标定图片的文件夹
chessboard_size = (8, 11)       # 棋盘格内角点数量 (宽度, 高度) = (列数-1, 行数-1)
square_size = 10.0              # 每个方格的实际尺寸（毫米）

# 准备物体点：类似 (0,0,0), (1,0,0), (2,0,0) ...., (7,10,0)
objp = np.zeros((chessboard_size[0] * chessboard_size[1], 3), np.float32)
objp[:, :2] = np.mgrid[0:chessboard_size[0], 0:chessboard_size[1]].T.reshape(-1, 2)
objp *= square_size  # 乘以实际尺寸

# 存储所有图像的对象点和图像点
objpoints = []  # 真实3D点
imgpoints = []  # 图像中的2D点

# 获取所有标定图片的文件列表
images = [f for f in os.listdir(chessboard_path) if f.endswith('.jpg')]
images.sort()  # 按文件名排序

if len(images) == 0:
    print(f"在目录 '{chessboard_path}' 中未找到jpg图片！")
    exit()

print(f"找到 {len(images)} 张标定图片")

# 遍历所有图片，检测角点
for i, fname in enumerate(images):
    img_path = os.path.join(chessboard_path, fname)
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图片: {fname}")
        continue
        
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # 查找棋盘格角点
    ret, corners = cv2.findChessboardCorners(gray, chessboard_size, None)
    
    if ret:
        # 精细化角点坐标
        criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
        corners_refined = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
        
        # 添加到数据集中
        objpoints.append(objp)
        imgpoints.append(corners_refined)
        
        # 在图像上绘制角点（可选，用于可视化）
        img_with_corners = cv2.drawChessboardCorners(img.copy(), chessboard_size, corners_refined, ret)
        
        # 显示处理进度
        cv2.putText(img_with_corners, f"Image {i+1}/{len(images)}", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        cv2.imshow('Corner Detection', img_with_corners)
        cv2.waitKey(300)  # 显示300ms
        
        print(f"图片 {fname}: 成功检测到角点")
    else:
        print(f"图片 {fname}: 未检测到角点")

cv2.destroyAllWindows()

# 检查是否有足够的有效图片
if len(objpoints) < 5:
    print(f"错误：只有 {len(objpoints)} 张图片成功检测到角点，至少需要5张！")
    exit()

print(f"\n开始相机标定，使用 {len(objpoints)} 张有效图片...")

# 进行相机标定
ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(
    objpoints, imgpoints, gray.shape[::-1], None, None
)

# 计算重投影误差（评估标定质量）
mean_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], camera_matrix, dist_coeffs)
    error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
    mean_error += error

mean_error /= len(objpoints)

print("\n=== 相机标定结果 ===")
print(f"重投影误差: {mean_error:.6f} 像素")
print("(值越小越好，通常应小于0.5像素)")

print(f"\n相机内参矩阵:")
print(camera_matrix)

print(f"\n畸变系数 (k1, k2, p1, p2, k3):")
print(dist_coeffs.flatten())

# 保存标定参数到YAML文件
# 修改标定程序中的YAML保存部分
# 替换原来的保存代码：

# 保存标定参数到YAML文件（修复版本）
calibration_data = {
    'calibration_date': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
    'image_count': len(objpoints),
    'reprojection_error': float(mean_error),
    'camera_matrix': camera_matrix.tolist(),
    'distortion_coefficients': dist_coeffs.flatten().tolist(),
    'chessboard_size': [chessboard_size[0], chessboard_size[1]],  # 改为列表而不是元组
    'square_size_mm': square_size,
    'image_resolution': [gray.shape[1], gray.shape[0]]  # 改为列表 [width, height]
}

# 使用安全的YAML转换
def represent_list(dumper, data):
    return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=True)

yaml.add_representer(list, represent_list)

with open('camera_calibration1280.yaml', 'w') as f:
    yaml.dump(calibration_data, f, default_flow_style=False, sort_keys=False)


# 可选：显示一张图片的矫正效果作为验证
if len(objpoints) > 0:
    # 使用第一张图片进行演示
    test_img_path = os.path.join(chessboard_path, images[0])
    test_img = cv2.imread(test_img_path)
    
    # 矫正畸变
    h, w = test_img.shape[:2]
    new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(
        camera_matrix, dist_coeffs, (w, h), 1, (w, h)
    )
    undistorted_img = cv2.undistort(test_img, camera_matrix, dist_coeffs, None, new_camera_matrix)
    
    # 裁剪图像
    x, y, w, h = roi
    undistorted_img = undistorted_img[y:y+h, x:x+w]
    
    # 并排显示原图和矫正后的图
    comparison = np.hstack((test_img, undistorted_img))
    comparison = cv2.resize(comparison, (1200, 400))
    
    cv2.putText(comparison, "Original (with distortion)", (50, 30), 
               cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    cv2.putText(comparison, "Undistorted", (650, 30), 
               cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
    
    cv2.imshow('Original vs Undistorted (Press any key to exit)', comparison)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

print("\n=== 标定完成 ===")
print("您可以使用保存的 camera_calibration.yaml 文件在后续程序中进行图像矫正。")